from querier.esquerier import ElasticSearchQuerier


class WeiboUserEngagementQuerier(ElasticSearchQuerier):

    def __init__(self, es, index, doc_type, nlp_service=None):
        super(WeiboUserEngagementQuerier, self).__init__(None, None, None)
        self.es = es
        self.index = index
        self.doc_type = doc_type
        self.nlp_service = nlp_service

    def _build_query(self, args):
        ids = args.get('user_ids')
        size = args.get('size')
        if size is None:
            size = 1000
        if ids is None:
            raise ValueError('message: "user_ids"(list of "user_id") is needed')

        query = {"query": {"terms": {"user_id": ids}}, "size": size}
        return query, {}, {'user_ids': ids}

    def _build_result(self, es_result, param):
        # ids = param['user_ids']
        users = []
        for hit in es_result['hits']['hits']:
            data = self.extract_result(hit)
            users.append(data)
        user_dict = {}
        for u in users:
            user_id = u['user_id']
            user_dict[user_id] = user_dict.get(user_id, []) + [u]
        return user_dict

    @staticmethod
    def extract_result(hit):
        source_ = hit['_source']
        res = {
            'user_id': source_['user_id'],
            'year_month': source_['year_month'],
            "sum_retweets": source_['sum_retweets'],
            "avg_retweets": source_['avg_retweets'],
            "max_retweets": source_['max_retweets'],
            "min_retweets": source_['min_retweets'],

            "sum_comments": source_['sum_comments'],
            "avg_comments": source_['avg_comments'],
            "max_comments": source_['max_comments'],
            "min_comments": source_['min_comments'],

            "sum_likes": source_['sum_likes'],
            "avg_likes": source_['avg_likes'],
            "max_likes": source_['max_likes'],
            "min_likes": source_['min_likes'],

            "sum_engagement": source_['sum_engagement'],
            "avg_engagement": source_['avg_engagement'],
            "max_engagement": source_['max_engagement'],
            "min_engagement": source_['min_engagement'],
            "count": source_['count'],
        }

        return res
